#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF广告内容删除工具
功能：
1. 删除指定的文本广告内容
2. 删除水印
3. 删除固定尺寸的图片
4. 批量处理PDF文件
"""

import os
import fitz  # PyMuPDF
import re
import logging
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Union
import shutil
import sys

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class PDFAdRemover:
    """PDF广告内容删除器"""
    
    def __init__(self):
        self.processed_files = 0
        self.removed_ads = 0
        self.errors = 0
    
    def _handle_chinese_path(self, file_path: str) -> str:
        """
        处理中文路径问题，创建临时符号链接或复制文件
        
        Args:
            file_path (str): 原始文件路径
            
        Returns:
            str: 处理后的文件路径
        """
        try:
            # 检查路径是否包含中文字符
            if any('\u4e00' <= char <= '\u9fff' for char in file_path):
                logger.info(f"检测到中文路径: {file_path}")
                
                # 创建临时目录
                temp_dir = os.path.join(os.path.dirname(__file__), "temp")
                os.makedirs(temp_dir, exist_ok=True)
                
                # 生成临时文件名
                import uuid
                temp_filename = f"temp_{uuid.uuid4().hex}.pdf"
                temp_path = os.path.join(temp_dir, temp_filename)
                
                # 复制文件到临时路径
                shutil.copy2(file_path, temp_path)
                logger.info(f"已复制到临时路径: {temp_path}")
                
                return temp_path
            else:
                return file_path
        except Exception as e:
            logger.error(f"处理中文路径时出错: {str(e)}")
            return file_path
    
    def _process_pdf_with_temp_copy(self, pdf_path: str, process_func):
        """
        使用临时文件处理PDF，解决中文路径问题
        
        Args:
            pdf_path (str): 原始PDF路径
            process_func: 处理函数，接受处理后的路径和输出路径
            
        Returns:
            bool: 是否成功处理
        """
        temp_path = None
        try:
            # 检查是否需要处理中文路径
            if any('\u4e00' <= char <= '\u9fff' for char in pdf_path):
                logger.info(f"检测到中文路径，创建临时副本: {pdf_path}")
                
                # 创建临时目录
                temp_dir = os.path.join(os.path.dirname(__file__), "temp")
                os.makedirs(temp_dir, exist_ok=True)
                
                # 生成临时文件名
                import uuid
                temp_filename = f"temp_{uuid.uuid4().hex}.pdf"
                temp_path = os.path.join(temp_dir, temp_filename)
                
                # 复制文件到临时路径
                shutil.copy2(pdf_path, temp_path)
                logger.info(f"已复制到临时路径: {temp_path}")
                
                # 调用处理函数
                result = process_func(temp_path)
                
                if result:
                    # 将处理后的文件复制回原位置
                    shutil.copy2(temp_path, pdf_path)
                    logger.info(f"已将处理结果复制回原位置: {pdf_path}")
                
                return result
            else:
                # 直接处理
                return process_func(pdf_path)
                
        except Exception as e:
            logger.error(f"处理PDF时出错: {str(e)}")
            return False
        finally:
            # 清理临时文件
            if temp_path and os.path.exists(temp_path):
                try:
                    os.remove(temp_path)
                except:
                    pass
    
    def remove_text_ads(self, pdf_path: str, 
                       ad_texts: List[str], 
                       output_path: Optional[str] = None,
                       case_sensitive: bool = False) -> bool:
        """
        删除PDF中的指定文本广告内容
        
        Args:
            pdf_path (str): PDF文件路径
            ad_texts (List[str]): 要删除的广告文本列表
            output_path (str, optional): 输出文件路径，默认覆盖原文件
            case_sensitive (bool): 是否区分大小写
            
        Returns:
            bool: 是否成功处理
        """
        def _process_text_ads(processed_path):
            """内部处理函数"""
            try:
                # 打开PDF文件
                doc = fitz.open(processed_path)
                modified = False
                
                # 遍历所有页面
                for page_num in range(len(doc)):
                    page = doc[page_num]
                    
                    # 方法1: 处理跨span的长文本广告
                    full_text = page.get_text()
                    for ad_text in ad_texts:
                        if self._text_contains_ad(full_text, ad_text, case_sensitive):
                            # 找到广告文本在页面中的位置
                            removed_spans = self._remove_long_text_ad(page, ad_text, case_sensitive)
                            if removed_spans > 0:
                                modified = True
                                self.removed_ads += 1
                                logger.info(f"页面 {page_num + 1}: 删除长文本广告 '{ad_text[:50]}...' ({removed_spans} 个文本块)")
                    
                    # 方法2: 处理单个span内的广告文本
                    text_dict = page.get_text("dict")
                    
                    # 遍历每个文本块
                    for block in text_dict["blocks"]:
                        if "lines" in block:  # 文本块
                            for line in block["lines"]:
                                for span in line["spans"]:
                                    text = span["text"]
                                    
                                    # 检查是否包含广告文本
                                    for ad_text in ad_texts:
                                        if self._text_contains_ad(text, ad_text, case_sensitive):
                                            # 创建白色矩形覆盖广告文本
                                            bbox = span["bbox"]
                                            rect = fitz.Rect(bbox)
                                            
                                            # 添加白色矩形覆盖
                                            page.draw_rect(rect, color=fitz.utils.getColor("white"), 
                                                         fill=fitz.utils.getColor("white"))
                                            modified = True
                                            self.removed_ads += 1
                                            logger.info(f"页面 {page_num + 1}: 删除广告文本 '{ad_text[:50]}...'")
                
                if modified:
                    # 保存文件
                    doc.save(processed_path)
                    logger.info(f"成功处理文件: {os.path.basename(pdf_path)}")
                
                doc.close()
                return True
                
            except Exception as e:
                logger.error(f"处理文件 {pdf_path} 时出错: {str(e)}")
                self.errors += 1
                return False
        
        # 使用临时文件处理
        return self._process_pdf_with_temp_copy(pdf_path, _process_text_ads)
    
    def remove_watermarks(self, pdf_path: str, 
                         watermark_patterns: List[str] = None,
                         output_path: Optional[str] = None) -> bool:
        """
        删除PDF中的水印
        
        Args:
            pdf_path (str): PDF文件路径
            watermark_patterns (List[str]): 水印文本模式列表，默认包含常见水印
            output_path (str, optional): 输出文件路径，默认覆盖原文件
            
        Returns:
            bool: 是否成功处理
        """
        if watermark_patterns is None:
            watermark_patterns = [
                "水印", "watermark", "机密", "confidential", "草稿", "draft",
                "样本", "sample", "内部", "internal", "禁止复制", "copyright",
                "版权所有", "版权", "©", "®", "™"
            ]
        
        try:
            doc = fitz.open(pdf_path)
            modified = False
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                
                # 获取页面尺寸
                page_rect = page.rect
                
                # 获取文本块
                text_dict = page.get_text("dict")
                
                for block in text_dict["blocks"]:
                    if "lines" in block:
                        for line in block["lines"]:
                            for span in line["spans"]:
                                text = span["text"].strip()
                                
                                # 检查是否为水印
                                if self._is_watermark(text, watermark_patterns):
                                    bbox = span["bbox"]
                                    rect = fitz.Rect(bbox)
                                    
                                    # 创建半透明白色矩形覆盖水印
                                    page.draw_rect(rect, color=fitz.utils.getColor("white"), 
                                                 fill=fitz.utils.getColor("white"))
                                    modified = True
                                    self.removed_ads += 1
                                    logger.info(f"页面 {page_num + 1}: 删除水印 '{text}'")
                
                # 检查页面底部的常见水印位置
                self._remove_bottom_watermarks(page, page_num, watermark_patterns)
            
            if modified:
                if output_path is None:
                    output_path = pdf_path
                doc.save(output_path)
                logger.info(f"成功处理文件: {os.path.basename(pdf_path)}")
            
            doc.close()
            return True
            
        except Exception as e:
            logger.error(f"处理文件 {pdf_path} 时出错: {str(e)}")
            self.errors += 1
            return False
    
    def remove_images_by_size(self, pdf_path: str, 
                             max_width: float = 100,
                             max_height: float = 100,
                             output_path: Optional[str] = None) -> bool:
        """
        删除指定尺寸范围内的图片（通常是小广告图片）
        
        Args:
            pdf_path (str): PDF文件路径
            max_width (float): 最大宽度（点）
            max_height (float): 最大高度（点）
            output_path (str, optional): 输出文件路径，默认覆盖原文件
            
        Returns:
            bool: 是否成功处理
        """
        try:
            doc = fitz.open(pdf_path)
            modified = False
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                
                # 获取页面上的图片列表
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    # 获取图片在页面上的位置
                    img_rects = page.get_image_rects(img[0])
                    
                    for rect in img_rects:
                        width = rect.width
                        height = rect.height
                        
                        # 检查图片尺寸是否在指定范围内
                        if width <= max_width and height <= max_height:
                            # 用白色矩形覆盖小图片
                            page.draw_rect(rect, color=fitz.utils.getColor("white"), 
                                         fill=fitz.utils.getColor("white"))
                            modified = True
                            self.removed_ads += 1
                            logger.info(f"页面 {page_num + 1}: 删除小图片 ({width:.1f}x{height:.1f})")
            
            if modified:
                if output_path is None:
                    output_path = pdf_path
                doc.save(output_path)
                logger.info(f"成功处理文件: {os.path.basename(pdf_path)}")
            
            doc.close()
            return True
            
        except Exception as e:
            logger.error(f"处理文件 {pdf_path} 时出错: {str(e)}")
            self.errors += 1
            return False
    
    def remove_ads_comprehensive(self, pdf_path: str,
                                ad_texts: List[str] = None,
                                watermark_patterns: List[str] = None,
                                max_image_width: float = 100,
                                max_image_height: float = 100,
                                output_path: Optional[str] = None,
                                create_backup: bool = True) -> Dict[str, int]:
        """
        综合删除PDF中的各种广告内容
        
        Args:
            pdf_path (str): PDF文件路径
            ad_texts (List[str]): 要删除的广告文本列表
            watermark_patterns (List[str]): 水印模式列表
            max_image_width (float): 最大图片宽度
            max_image_height (float): 最大图片高度
            output_path (str, optional): 输出文件路径
            create_backup (bool): 是否创建备份文件
            
        Returns:
            Dict[str, int]: 处理结果统计
        """
        results = {
            "processed": 0,
            "text_ads_removed": 0,
            "watermarks_removed": 0,
            "images_removed": 0,
            "errors": 0
        }
        
        try:
            # 创建备份
            if create_backup and output_path is None:
                backup_path = self._create_backup(pdf_path)
                logger.info(f"创建备份文件: {backup_path}")
            
            # 处理文本广告
            if ad_texts:
                if self.remove_text_ads(pdf_path, ad_texts, output_path):
                    results["text_ads_removed"] = self.removed_ads
            
            # 处理水印
            if watermark_patterns:
                if self.remove_watermarks(pdf_path, watermark_patterns, output_path):
                    results["watermarks_removed"] = self.removed_ads
            
            # 处理小图片
            if self.remove_images_by_size(pdf_path, max_image_width, max_image_height, output_path):
                results["images_removed"] = self.removed_ads
            
            results["processed"] = 1
            self.processed_files += 1
            
        except Exception as e:
            logger.error(f"综合处理文件 {pdf_path} 时出错: {str(e)}")
            results["errors"] = 1
            self.errors += 1
        
        return results
    
    def batch_process_directory(self, directory_path: str,
                               ad_texts: List[str] = None,
                               watermark_patterns: List[str] = None,
                               max_image_width: float = 100,
                               max_image_height: float = 100,
                               output_directory: Optional[str] = None,
                               create_backup: bool = True) -> Dict[str, int]:
        """
        批量处理目录下的所有PDF文件
        
        Args:
            directory_path (str): 目录路径
            ad_texts (List[str]): 广告文本列表
            watermark_patterns (List[str]): 水印模式列表
            max_image_width (float): 最大图片宽度
            max_image_height (float): 最大图片高度
            output_directory (str, optional): 输出目录
            create_backup (bool): 是否创建备份
            
        Returns:
            Dict[str, int]: 处理结果统计
        """
        results = {
            "total_files": 0,
            "processed_files": 0,
            "text_ads_removed": 0,
            "watermarks_removed": 0,
            "images_removed": 0,
            "errors": 0
        }
        
        # 查找所有PDF文件
        pdf_files = self._find_pdf_files(directory_path)
        results["total_files"] = len(pdf_files)
        
        if not pdf_files:
            logger.warning(f"在目录 {directory_path} 中未找到PDF文件")
            return results
        
        # 创建输出目录
        if output_directory:
            os.makedirs(output_directory, exist_ok=True)
        
        # 处理每个PDF文件
        for i, pdf_path in enumerate(pdf_files, 1):
            logger.info(f"正在处理 ({i}/{len(pdf_files)}): {os.path.basename(pdf_path)}")
            
            # 确定输出路径
            if output_directory:
                output_path = os.path.join(output_directory, os.path.basename(pdf_path))
            else:
                output_path = None
            
            # 处理文件
            file_results = self.remove_ads_comprehensive(
                pdf_path, ad_texts, watermark_patterns,
                max_image_width, max_image_height, output_path, create_backup
            )
            
            # 累加结果
            for key in results:
                if key != "total_files":
                    results[key] += file_results.get(key, 0)
        
        return results
    
    def _text_contains_ad(self, text: str, ad_text: str, case_sensitive: bool) -> bool:
        """检查文本是否包含广告内容"""
        if not case_sensitive:
            text = text.lower()
            ad_text = ad_text.lower()
        return ad_text in text
    
    def _remove_long_text_ad(self, page, ad_text: str, case_sensitive: bool) -> int:
        """
        删除跨多个span的长文本广告
        
        Args:
            page: PDF页面对象
            ad_text (str): 要删除的广告文本
            case_sensitive (bool): 是否区分大小写
            
        Returns:
            int: 删除的文本块数量
        """
        removed_count = 0
        try:
            # 获取页面文本字典
            text_dict = page.get_text("dict")
            
            # 获取页面的完整文本
            full_text = page.get_text()
            
            # 找到广告文本在完整文本中的位置
            if not case_sensitive:
                search_text = full_text.lower()
                search_ad = ad_text.lower()
            else:
                search_text = full_text
                search_ad = ad_text
            
            if search_ad not in search_text:
                return 0
            
            # 计算广告文本在完整文本中的字符位置
            start_pos = search_text.find(search_ad)
            end_pos = start_pos + len(search_ad)
            
            # 遍历文本块，找到包含广告文本的span
            current_pos = 0
            for block in text_dict["blocks"]:
                if "lines" in block:  # 文本块
                    for line in block["lines"]:
                        for span in line["spans"]:
                            span_text = span["text"]
                            span_length = len(span_text)
                            
                            # 检查当前span是否与广告文本重叠
                            span_start = current_pos
                            span_end = current_pos + span_length
                            
                            # 如果span与广告文本有重叠，则删除
                            if (span_start < end_pos and span_end > start_pos):
                                bbox = span["bbox"]
                                rect = fitz.Rect(bbox)
                                
                                # 添加白色矩形覆盖
                                page.draw_rect(rect, color=fitz.utils.getColor("white"), 
                                             fill=fitz.utils.getColor("white"))
                                removed_count += 1
                            
                            current_pos += span_length
            
            return removed_count
            
        except Exception as e:
            logger.warning(f"删除长文本广告时出错: {str(e)}")
            return 0
    
    def _is_watermark(self, text: str, patterns: List[str]) -> bool:
        """检查文本是否为水印"""
        text_lower = text.lower()
        for pattern in patterns:
            if pattern.lower() in text_lower:
                return True
        return False
    
    def _remove_bottom_watermarks(self, page, page_num: int, patterns: List[str]):
        """删除页面底部的水印"""
        try:
            # 获取页面底部区域的文本
            bottom_rect = fitz.Rect(0, page.rect.height - 50, page.rect.width, page.rect.height)
            text_dict = page.get_text("dict", clip=bottom_rect)
            
            for block in text_dict["blocks"]:
                if "lines" in block:
                    for line in block["lines"]:
                        for span in line["spans"]:
                            text = span["text"].strip()
                            if self._is_watermark(text, patterns):
                                bbox = span["bbox"]
                                rect = fitz.Rect(bbox)
                                page.draw_rect(rect, color=fitz.utils.getColor("white"), 
                                             fill=fitz.utils.getColor("white"))
                                logger.info(f"页面 {page_num + 1}: 删除底部水印 '{text}'")
        except Exception as e:
            logger.warning(f"处理页面 {page_num + 1} 底部水印时出错: {str(e)}")
    
    def _create_backup(self, pdf_path: str) -> str:
        """创建PDF文件的备份"""
        backup_path = pdf_path.replace('.pdf', '_backup.pdf')
        shutil.copy2(pdf_path, backup_path)
        return backup_path
    
    def _find_pdf_files(self, directory_path: str) -> List[str]:
        """查找目录下的所有PDF文件"""
        pdf_files = []
        if not os.path.exists(directory_path):
            logger.error(f"目录不存在: {directory_path}")
            return pdf_files
        
        for root, dirs, files in os.walk(directory_path):
            for file in files:
                if file.lower().endswith('.pdf'):
                    pdf_path = os.path.join(root, file)
                    pdf_files.append(pdf_path)
        
        return pdf_files


def main():
    """主函数 - 提供命令行接口"""
    import argparse
    
    parser = argparse.ArgumentParser(description="PDF广告内容删除工具")
    parser.add_argument("input", help="输入PDF文件或目录路径")
    parser.add_argument("--output", help="输出文件或目录路径")
    parser.add_argument("--ad-texts", nargs="+", help="要删除的广告文本列表")
    parser.add_argument("--watermarks", nargs="+", help="水印模式列表")
    parser.add_argument("--max-width", type=float, default=100, help="最大图片宽度")
    parser.add_argument("--max-height", type=float, default=100, help="最大图片高度")
    parser.add_argument("--no-backup", action="store_true", help="不创建备份文件")
    
    args = parser.parse_args()
    
    # 创建广告删除器
    remover = PDFAdRemover()
    
    # 默认广告文本和水印模式
    ad_texts = args.ad_texts or [
        "广告", "推广", "优惠", "折扣", "促销", "特价",
        "立即购买", "点击这里", "了解更多", "免费试用"
    ]
    
    watermark_patterns = args.watermarks or [
        "水印", "watermark", "机密", "confidential", "草稿", "draft",
        "样本", "sample", "内部", "internal", "禁止复制", "copyright"
    ]
    
    print("PDF广告内容删除工具")
    print(f"输入路径: {args.input}")
    print(f"输出路径: {args.output or '覆盖原文件'}")
    print(f"广告文本: {ad_texts}")
    print(f"水印模式: {watermark_patterns}")
    print(f"最大图片尺寸: {args.max_width}x{args.max_height}")
    print(f"创建备份: {'否' if args.no_backup else '是'}")
    print("-" * 60)
    
    # 处理命令行编码问题
    try:
        # 尝试修复编码问题
        if sys.platform == "win32":
            # Windows系统可能需要特殊处理
            import locale
            system_encoding = locale.getpreferredencoding()
            print(f"系统编码: {system_encoding}")
            
            # 尝试不同的编码方式
            try:
                # 先尝试用系统编码解码再编码
                fixed_input = args.input.encode(system_encoding).decode('utf-8')
                input_path = os.path.abspath(fixed_input)
                print(f"修复后路径: {input_path}")
            except:
                # 如果失败，使用原始路径
                input_path = os.path.abspath(args.input)
                print(f"使用原始路径: {input_path}")
        else:
            input_path = os.path.abspath(args.input)
    except Exception as e:
        print(f"路径处理错误: {e}")
        input_path = os.path.abspath(args.input)
    
    print(f"最终路径: {input_path}")
    print(f"路径存在: {os.path.exists(input_path)}")
    print(f"是文件: {os.path.isfile(input_path)}")
    print(f"是目录: {os.path.isdir(input_path)}")
    
    # 判断是文件还是目录
    if os.path.isfile(input_path):
        # 处理单个文件
        results = remover.remove_ads_comprehensive(
            input_path, ad_texts, watermark_patterns,
            args.max_width, args.max_height, args.output, not args.no_backup
        )
        print(f"处理完成！删除了 {results.get('text_ads_removed', 0)} 个文本广告")
        print(f"删除了 {results.get('watermarks_removed', 0)} 个水印")
        print(f"删除了 {results.get('images_removed', 0)} 个小图片")
    elif os.path.isdir(input_path):
        # 处理目录
        results = remover.batch_process_directory(
            input_path, ad_texts, watermark_patterns,
            args.max_width, args.max_height, args.output, not args.no_backup
        )
        print(f"批量处理完成！")
        print(f"总文件数: {results['total_files']}")
        print(f"已处理: {results['processed_files']}")
        print(f"文本广告: {results['text_ads_removed']}")
        print(f"水印: {results['watermarks_removed']}")
        print(f"小图片: {results['images_removed']}")
        print(f"错误: {results['errors']}")
    else:
        print(f"错误: 路径不存在 {input_path}")


if __name__ == "__main__":
    main()
